import torch
import torch.nn as nn

class ChannelAttention(nn.Module):
    """通道注意力模块
    专门用于增强低对比度图像中文字特征的通道注意力机制
    
    网络结构：
    --------------------------------
    输入特征图 [B, C, H, W]
          ↓
    并行双路注意力：
        - 全局平均池化 [B, C, 1, 1]
        - 全局最大池化 [B, C, 1, 1]
          ↓
    共享的全连接层：
        FC1(C→C/r) → ReLU → FC2(C/r→C) → Sigmoid
          ↓
    双路注意力权重相加 → Sigmoid
          ↓
    与输入特征图逐通道相乘 [B, C, H, W]
    --------------------------------
    
    作用原理：
    1. 通过分析每个通道的全局信息，自动学习通道重要性权重
    2. 增强文字相关通道的响应，抑制背景噪声通道
    3. 特别适用于背景与文字颜色接近的低对比度场景
    
    参数：
        in_channels (int): 输入特征图的通道数
        reduction (int): 特征压缩比率(default=8)
    """
    def __init__(self, in_channels, reduction=8):
        super(ChannelAttention, self).__init__()
        # 双路注意力池化层
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局平均池化：提取通道均值特征
        self.max_pool = nn.AdaptiveMaxPool2d(1)  # 全局最大池化：提取通道显著特征
        
        # 共享的注意力学习网络
        self.fc = nn.Sequential(
            # 第一阶段：降维压缩
            nn.Linear(in_channels, in_channels // reduction),  # [C] → [C/r]
            nn.ReLU(inplace=True),  # 引入非线性
            # 第二阶段：恢复维度
            nn.Linear(in_channels // reduction, in_channels),  # [C/r] → [C]
            nn.Sigmoid()  # 输出0-1的归一化权重
        )

    def forward(self, x):
        """前向传播
        
        参数：
            x: 输入特征图，形状 [B, C, H, W]
        
        返回：
            加权后的特征图，形状保持 [B, C, H, W]
        """
        batch_size, channels, _, _ = x.size()
        
        # 平均池化路径
        avg_out = self.fc(self.avg_pool(x).view(batch_size, channels))  # [B,C,1,1]→[B,C]→权重
        
        # 最大池化路径 
        max_out = self.fc(self.max_pool(x).view(batch_size, channels))  # 同上
        
        # 合并双路注意力（元素相加后通过Sigmoid）
        channel_weights = torch.sigmoid(avg_out + max_out)  # [B,C]
        
        # 将学习到的权重应用于输入特征图
        return x * channel_weights.view(batch_size, channels, 1, 1)  # 广播乘法

class CnnNet(nn.Module):
    """用于手写数字识别的CNN网络（输入28x28图像，输出10分类）
    
    网络结构：
    Conv1(1->32,5x5) -> BN -> ReLU -> MaxPool -> 
    Conv2(32->64,3x3) -> BN -> ReLU -> MaxPool -> 
    Conv3(64->128,3x3) -> BN -> ReLU -> MaxPool -> 
    Flatten -> FC1(自动计算->256) -> ReLU -> FC2(256->10)"""

    def __init__(self):
        super().__init__()
        # 卷积层定义
        # 第一层使用较大kernel捕捉网格特征
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)  # 输入1通道，输出32通道，3x3卷积核，步长1
        self.bn1 = nn.BatchNorm2d(32)
        # 第二层使用小kernel捕捉细节
        self.conv2 = nn.Conv2d(32, 64, 3, 1) # 输入32通道，输出64通道
        self.bn2 = nn.BatchNorm2d(64)
        # 增加第三卷积层
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool = nn.MaxPool2d(2, 2)
        # self.ca = ChannelAttention(128)
        # Dropout层
        self.dropout1 = nn.Dropout(0.25) # 25%的神经元失活
        self.dropout2 = nn.Dropout(0.5) # 50%的神经元失活
        # 自动计算全连接层输入尺寸
        self._to_linear = None
        self._calculate_conv_output((1, 28, 28))  # 默认MNIST输入尺寸
        self.fc1 = nn.Linear(self._to_linear, 256)  # 调整全连接层大小
        #self.fc1 = nn.Linear(128 * 3 * 3, 256)  # 调整全连接层大小
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.bn1(self.conv1(x))))
        x = self.pool(torch.relu(self.bn2(self.conv2(x))))
        x = self.pool(torch.relu(self.bn3(self.conv3(x))))
        # x = self.ca(x)  # 应用通道注意力
        x = self.dropout1(x)
        x = x.view(-1, self._to_linear)  # 使用自动计算的维度
        
        x = torch.relu(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        return x
    
    def _calculate_conv_output(self, shape):
        """自动计算卷积层输出后的展平尺寸"""
        with torch.no_grad():  # 不计算梯度
            # 创建虚拟输入(1, C, H, W)
            dummy = torch.zeros(1, *shape)
            
            # 模拟前向传播过程
            dummy = self.pool(torch.relu(self.bn1(self.conv1(dummy))))
            dummy = self.pool(torch.relu(self.bn2(self.conv2(dummy))))
            dummy = self.pool(torch.relu(self.bn3(self.conv3(dummy))))
            
            # 计算展平后的维度(排除batch维度)
            self._to_linear = dummy.view(1, -1).size(1)
            print(f"自动计算全连接层输入尺寸: {self._to_linear}")